/*
** Copyright (C) 2010 Calin Radoni
**
** This file is part of CHScanner.
**
** CHScanner is free software: you can redistribute it and/or modify
** it under the terms of the GNU General Public License as published by
** the Free Software Foundation, either version 3 of the License, or
** (at your option) any later version.
**
** CHScanner is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
** GNU General Public License for more details.
**
** You should have received a copy of the GNU General Public License
** along with CHScanner. If not, see <http://www.gnu.org/licenses/>.
*/

#include "StdAfx.h"
#include "ScanMethodDHCP.h"

//--------------------------------------------------------------------

ScanMethodDHCP::ScanMethodDHCP(void)
{
	myType = __wmtDhcpDiscov;
}

ScanMethodDHCP::~ScanMethodDHCP(void)
{
}

//--------------------------------------------------------------------

void ScanMethodDHCP::GenerateRandomMAC(void)
{
	BYTE randMacAddr[6];
	in_addr ia;

	ia.S_un.S_addr = GetRandomValue();

	randMacAddr[0] = 0;
	randMacAddr[1] = 0;
	randMacAddr[2] = ia.S_un.S_un_b.s_b1;
	randMacAddr[3] = ia.S_un.S_un_b.s_b2;
	randMacAddr[4] = ia.S_un.S_un_b.s_b3;
	randMacAddr[5] = ia.S_un.S_un_b.s_b4;

	macAddr.LoadFromBytePtr(randMacAddr);
}

char ScanMethodDHCP::GetCharFromByte(BYTE bbb, bool alfaOnly)
{
	BYTE val;

	val = bbb;

	if(alfaOnly){
		while(val>51)
			val -= 51;

		if(val<26){
			val += 'A';
			return (char)val;
		}

		val -= 26;
		val += 'a';
		return (char)val;
	}

	while(val>62)
		val -= 62;

	if(val==0){
		val = '-';
		return (char)val;
	}

	val--;
	if(val<10){
		val += '0';
		return (char)val;
	}

	val -= 10;
	if(val<26){
		val += 'A';
		return (char)val;
	}

	val -= 26;
	val += 'a';
	return (char)val;
}

void ScanMethodDHCP::GenerateRandomName(void)
{
	BYTE cnt, idx;
	BYTE randName[16];
	DWORD dwVal;

	dwVal = GetRandomValue();
	memcpy(&randName[0], &dwVal, 4);
	dwVal = GetRandomValue();
	memcpy(&randName[4], &dwVal, 4);
	dwVal = GetRandomValue();
	memcpy(&randName[8], &dwVal, 4);
	dwVal = GetRandomValue();
	memcpy(&randName[12], &dwVal, 4);

	cnt = randName[15];
	while(cnt>10)
		cnt -= 10;
	cnt += 5;

	for(idx=0; idx<cnt; idx++)
		if(idx==0)
			randName[idx] = GetCharFromByte(randName[idx], true);
		else
			randName[idx] = GetCharFromByte(randName[idx], false);

	randName[cnt] = 0;

	reqName = (char*)randName;
}

//--------------------------------------------------------------------

int ScanMethodDHCP::InternalInit(void)
{
	CString str, strMac, strReqName;
	int pos;

	FillRandomBuffer();
	
	if(workPhase!=__wptDiscovery) return MSR_ERR_PHASE;

	str = scanParameters.Trim();
	if(str.GetLength()==0){
		// no parameters specified
		strMac = "*";
		strReqName = "*";
	}
	else{
		// at least one parameter
		pos = str.Find(' ');
		if(pos==-1){
			// only one parameter specified
			strMac = str;
			strReqName = "*";
		}
		else{
			// both parameters are specified
			strMac = str.Left(pos);
			strReqName = str.Mid(pos+1);
		}
	}

	if(strMac.CompareNoCase("*")==0){
		// generate random MAC
		GenerateRandomMAC();
	}
	else{
		if(!macAddr.LoadFromString(strMac))
			return MSR_ERR_PARAMETER;
	}

	if(strReqName.CompareNoCase("*")==0){
		// generate a random name
		GenerateRandomName();
	}
	else{
		if(strReqName.GetLength()>40) reqName = strReqName.Left(40);
		else                          reqName = strReqName;
	}

	tempLeaseTime = 0;

	return MSR_ERR_OK;
}

int ScanMethodDHCP::InternalExecute(void)
{
	ScanPacket* pack;
	bool res;

	while(PacketsUsed()>0){
		CheckResponses();
		CheckDeaths();
		SleepEx(1, TRUE);
	}

	pack = FindFirstUnusedPacket();
	while(pack==NULL){
		SleepEx(waitTime, TRUE);
		CheckResponses();
		CheckDeaths();
		pack = FindFirstUnusedPacket();
	}

	if(pack!=NULL){
		pack->Initialize(true);
		pack->sTime       = GetTickCount();
		pack->addr        = 0;
		pack->port        = 68;
		pack->methodType  = myType;
		pack->fullAddress = NULL;
		pack->machineInfo = NULL;

		res = SendDHCPDiscoverPacket(pack);
		if(!res)
			pack->Initialize(false);
	}
	else res = false;

	while(PacketsUsed()>0){
		CheckResponses();
		CheckDeaths();
		SleepEx(1, TRUE);
	}

	if(hWnd!=NULL)
		PostMessage(hWnd, CHS_MSG_PACKETS_SENT, (WPARAM)0, (LPARAM)packSent);

	if(!res) return MSR_ERR_FAILED;

	return MSR_ERR_OK;
}

//--------------------------------------------------------------------

void ScanMethodDHCP::ComputeRandomValues(void)
{
	transactionId = GetRandomValue();
}

//--------------------------------------------------------------------

bool ScanMethodDHCP::SendDHCPDiscoverPacket(ScanPacket* packIn)
{
	in_addr ia;
	BYTE srcMAC[6];
	BYTE dstMAC[6];
	int i;

	if(packIn==NULL) return false;

	ComputeRandomValues();

	ia.S_un.S_un_b.s_b1 = 0;
	ia.S_un.S_un_b.s_b2 = 0;
	ia.S_un.S_un_b.s_b3 = 0;
	ia.S_un.S_un_b.s_b4 = 0;

	memcpy(srcMAC, macAddr.addr, 6);

	for(i=0; i<6; i++)
		dstMAC[i] = 0xFF;

	packIn->paramDW = transactionId;
	for(i=0; i<6; i++)
		packIn->paramIA.u.Byte[i] = srcMAC[i];

	if(!packetBuilderEx->BuildPacketDHCPDiscover(srcMAC, ia.S_un.S_addr, reqName, transactionId))
		return false;

	if(!packetBuilderEx->AddEthernetHeader(srcMAC, dstMAC, 0x0800))
		return false;

	return SendPacket(packetBuilderEx->GetPacket(), (int)packetBuilderEx->GetPacketSize());
}

bool ScanMethodDHCP::SendDHCPRequestPacket(ScanPacket* packIn, DWORD dhcpSrvIp, DWORD reqIP)
{
	BYTE srcMAC[6];
	BYTE dstMAC[6];
	int i;

	if(packIn==NULL) return false;
	if(dhcpSrvIp==0) return false;

	for(i=0; i<6; i++)
		srcMAC[i] = packIn->paramIA.u.Byte[i];

	for(i=0; i<6; i++)
		dstMAC[i] = 0xFF;

	transactionId = packIn->paramDW;

	if(!packetBuilderEx->BuildPacketDHCPRequest(srcMAC, dhcpSrvIp, reqIP, reqName, transactionId, tempLeaseTime))
		return false;

	if(!packetBuilderEx->AddEthernetHeader(srcMAC, dstMAC, 0x0800))
		return false;

	return SendPacket(packetBuilderEx->GetPacket(), (int)packetBuilderEx->GetPacketSize());
}

//--------------------------------------------------------------------

bool ScanMethodDHCP::ParseOptions(DHCPMessage* msg, int dataOffset, int sizeIn)
{
	int currentPos;
	BYTE optId, optLen, oIdx;
	BYTE *bbb;
	int idx;
	bool done;
	bool error;
	MDHCPOption optDHCP;

	currentPos = dataOffset + 244;	// 244 is 8(=sizeof(UDPHeaderPG)) + 236(=sizeof(DHCPMessage Header))

	idx = 0;

	// Check for magic cookie
	if(msg->options[idx++]!=0x63) return false;
	if(msg->options[idx++]!=0x82) return false;
	if(msg->options[idx++]!=0x53) return false;
	if(msg->options[idx++]!=0x63) return false;

	mach.Initialize();

	mach.infoDHCP.isEmpty = false;

	mach.infoDHCP.hLen   = msg->hLen;
	mach.infoDHCP.yiAddr = msg->yiAddr;
	mach.infoDHCP.siAddr = msg->siAddr;
	mach.infoDHCP.giAddr = msg->giAddr;
	memcpy(mach.infoDHCP.chAddr, msg->chAddr, 16);

	// parse options
	done  = false;
	error = false;
	while(!done && !error){
		optId = msg->options[idx++];
		currentPos++;

		if(optId!=0xFF){
			if(currentPos<sizeIn){
				optLen = msg->options[idx++];
				currentPos++;

				switch(optId){
					/*** length = 4 ***/
					case 1 :		// Subnet Mask
					case 2 :		// Time Offset
					case 28 :		// Broadcast Address Option
					case 32 :		// Router Solicitation Address Option
					case 35 :		// ARP Cache Timeout Option
					case 38 :		// TCP Keepalive Interval Option
					case 50 :		// Requested IP Address
					case 51 :		// IP Address Lease Time
					case 54 :		// Server Identifier
					case 58 :		// Renewal (T1) Time Value
					case 59 :		// Rebinding (T2) Time Value
						if(optLen!=4) error = true;

						if(optId==51){
							bbb = (BYTE*)&tempLeaseTime;
							bbb[0] = msg->options[idx+3];
							bbb[1] = msg->options[idx+2];
							bbb[2] = msg->options[idx+1];
							bbb[3] = msg->options[idx];
						}
						if(optId==54){
							tempSrvId.S_un.S_un_b.s_b1 = msg->options[idx];
							tempSrvId.S_un.S_un_b.s_b2 = msg->options[idx+1];
							tempSrvId.S_un.S_un_b.s_b3 = msg->options[idx+2];
							tempSrvId.S_un.S_un_b.s_b4 = msg->options[idx+3];
						}
						break;

					/*** length = 4*n ***/
					case 3 :		// Router Option
					case 4 :		// Time Server Option
					case 5 :		// Name Server Option
					case 6 :		// Domain Name Server Option
					case 7 :		// Log Server Option
					case 8 :		// Cookie Server Option
					case 9 :		// LPR Server Option
					case 10 :		// Impress Server Option
					case 11 :		// Resource Location Server Option
					case 41 :		// Network Information Servers Option
					case 42 :		// Network Time Protocol Servers Option
					case 44 :		// NetBIOS over TCP/IP Name Server Option
					case 45 :		// NetBIOS over TCP/IP Datagram Distribution Server Option
					case 48 :		// X Window System Font Server Option
					case 49 :		// X Window System Display Manager Option
					case 65 :		// Network Information Service+ Servers Option
					case 69 :		// Simple Mail Transport Protocol (SMTP) Server Option
					case 70 :		// Post Office Protocol (POP3) Server Option
					case 71 :		// Network News Transport Protocol (NNTP) Server Option
					case 72 :		// Default World Wide Web (WWW) Server Option
					case 73 :		// Default Finger Server Option
					case 74 :		// Default Internet Relay Chat (IRC) Server Option
					case 75 :		// StreetTalk Server Option
					case 76 :		// StreetTalk Directory Assistance (STDA) Server Option
						if((optLen%4)!=0) error = true;
						break;

					default:
						break;
				} /* switch */

				if(!error){
					if(optDHCP.Create(optId, optLen)){
						for(oIdx=0; oIdx<optLen; oIdx++){
							optDHCP.optionData[oIdx] = msg->options[idx++];
							currentPos++;
						}

						if(!mach.infoDHCP.optList.Add(&optDHCP))
							error = true;
					}
					else error = true;
				}

				if(idx>=312 || currentPos>=sizeIn) error = true;
			} /* currentPos passes over the sizeIn */
			else error = true;
		} /* optId is NOT EndOfOptions */
		else done = true;
	} /* while */
	
	return !error;
}

void ScanMethodDHCP::ProcessUdpEx(BYTE* dataIn, int sizeIn, BYTE* srcMAC)
{
	IPHeaderPG   *iph;
	UDPHeaderPG  *udph;
	int          dataOffset;
	ScanPacket   *pack;
	int          len;
	DHCPMessage  *msg;
	MFullAddress addr;
	in_addr      ia;
	int          i;

	iph = (IPHeaderPG*)dataIn;
	dataOffset  = (int)(iph->HLV & 0x0F);
	dataOffset *= 4;

	/*
		minimum DHCP packet size should be:
			  8(=sizeof(UDPHeaderPG)) +
			236(=sizeof(DHCPMessage Header)) +
			  4(=sizeof(Magic Cookie)) +
			  1(=at least the FF option)
			  = 249
	*/
	if(sizeIn<(dataOffset + 249)){
		// the size is too small for an DHCP packet
		return;
	}

	udph = (UDPHeaderPG*)&dataIn[dataOffset];
	if(ntohs(udph->SourcePort)!=67){
		// not a DHCP reply packet
		return;
	}
	if(ntohs(udph->DestinationPort)!=68){
		// not a DHCP reply packet
		return;
	}

	len = ntohs(udph->Length);
	if(sizeIn<(len+dataOffset)){
		// incomplete UDP packet
		return;
	}

	if(len<241){
		// incomplete DHCP packet
		return;
	}

	msg = (DHCPMessage*)&dataIn[dataOffset+8];

	if(msg->op!=2){
		// not a BOOTREPLY packet
		return;
	}

	pack = FindPacket(0, 68, myType);
	if(pack==NULL) return;

	if(msg->xid!=pack->paramDW){
		// not my transaction ID

		// I should memorize the address of this DHCP server anyway ? -> NO

		return;
	}

	tempSrvId.S_un.S_addr = 0;
	tempLeaseTime = 0;
	if(!ParseOptions(msg, dataOffset, sizeIn)){
		// incorrect option list or memory error
		return;
	}

	addr.Initialize();
	if(tempSrvId.S_un.S_addr!=0)
		addr.addrIP.Set(tempSrvId.S_un.S_addr);
	else
		addr.addrIP.Set(msg->siAddr);

	ia.S_un.S_un_b.s_b1 = iph->Source[0];
	ia.S_un.S_un_b.s_b2 = iph->Source[1];
	ia.S_un.S_un_b.s_b3 = iph->Source[2];
	ia.S_un.S_un_b.s_b4 = iph->Source[3];
	if(tempSrvId.S_un.S_addr==ia.S_un.S_addr){
		for(i=0; i<6; i++)
			addr.addrMAC.addr[i] = srcMAC[i];
		addr.gwMAC = false;
	}
	else addr.gwMAC = true;

	addr.alive = true;
	addr.respReceived = true;

	mach.address.LoadFromPointer(&addr);

	machInfoList->Add(&mach);
}

//--------------------------------------------------------------------
